package org.example.function;

import com.antgroup.geaflow.api.context.RuntimeContext;
import com.antgroup.geaflow.api.function.RichFunction;
import com.antgroup.geaflow.api.function.io.SourceFunction;
import com.antgroup.geaflow.api.window.IWindow;
import org.example.csv.CSVData;
import org.example.util.CSVFaster;

import java.io.Serializable;
import java.nio.file.Path;
import java.util.Collection;
import java.util.List;

public class CSVFileSourceV2<OUT> extends RichFunction implements SourceFunction<OUT> {

    private final Path filePath;

    private final List<String> selectedHeaders;

    private final CSVLineParser<OUT> parser;

    private int readPos = 0;

    private List<String[]> lines;

    public CSVFileSourceV2(Path filePath, List<String> selectedHeaders, CSVLineParser<OUT> parser) {
        this.filePath = filePath;
        this.selectedHeaders = selectedHeaders;
        this.parser = parser;
    }

    @Override
    public void open(RuntimeContext runtimeContext) {
    }

    @Override
    public void close() {

    }

    @Override
    public void init(int parallel, int index) {
        CSVData csvData = CSVFaster.read(this.filePath, selectedHeaders);
        this.lines = csvData.getLines();
    }

    @Override
    public boolean fetch(IWindow<OUT> window, SourceContext<OUT> ctx) throws Exception {
        while (readPos < this.lines.size()) {
            Collection<OUT> outs = this.parser.parse(this.lines.get(readPos));
            for (OUT out : outs) {
                ctx.collect(out);
            }
            readPos++;
        }

        return readPos <= this.lines.size();
    }

    public interface CSVLineParser<OUT> extends Serializable {

        Collection<OUT> parse(String[] line);
    }
}
